#include "unit.h"

// Defaut constructor returns an empty uninitialized unit.
Unit::Unit() : initialized(false), sort(&atoms), filter(&atoms,&offsets), grid(&atoms,&offsets) { }
// Defaut constructor returns an uninitialized unit with atoms alist and periodicity p.
Unit::Unit(const Atom::list &alist, const imat &p) : initialized(false),atoms(alist), sort(&atoms), filter(&atoms,&offsets), grid(&atoms,&offsets) { setPeriodicity(p); }
// Defaut constructor returns an uninitialized unit with atoms alist and periodicity p.
Unit::Unit(const std::vector<ivec4> &alist, const std::vector<ivec4> &p) : initialized(false), sort(&atoms), filter(&atoms,&offsets), grid(&atoms,&offsets)
{
	const size_t size = alist.size();
	atoms.reserve(size);
	for(size_t a = 0; a < size; a++)
		atoms.push_back(Atom(alist[a]));
		
	const size_t psize = p.size();
	imat per(4,psize);
	for(size_t i = 0; i < psize; i++)
		per(arma::span::all,i) = p[i];

	setPeriodicity(per);

	filter();	// Call filter object.
	sort();		// Call sorter object.
}
	
// This setter function replaces the current atom list with alist. Note this resets the unit, see reset().
void Unit::setAtoms(const Atom::list &alist) { atoms = alist; reset(); }

// This getter function return a const reference to the atoms. Read only.
const Atom::list &Unit::getAtoms() const { return atoms; }
// This getter function return a const reference to the atoms. Read only.
const Bond::table &Unit::getBonds() const { return bonds; }
// This getter function return a const reference to the atoms. Read only.
const umat &Unit::getBondsLocations() const { return bondlocations; }
	
// This getter function returns weither the unit has been initialized.
bool Unit::isInitilized() const { return initialized; }
// This getter function returns the periodicity matrix.
const imat &Unit::getPeriodicity() const { return periodicity; }
// This getter function returns the periodicity matrix.
const mat &Unit::getReciprocity() const { return reciprocity; }
// This getter function returns the periodicity matrix.
const mat &Unit::getBrillouin() const { return brillouin; }
// This getter function returns the periodicity matrix.
mat Unit::getPeriodicityXyz() const { return toXyz(periodicity); }

#include <functional>

imat Unit::recursive_find_p_vectors (const imat &per)
{
	arma::imat result(per.n_rows,pow(3,per.n_cols)-1);
	result.col(0) = per.col(0);
	result.col(1) = -per.col(0);

	if(per.n_cols != 1)
	{
		imat tmp1 = recursive_find_p_vectors(per.cols(1,per.n_cols-1));
		imat tmp2 = tmp1;
		result.cols(2,1+tmp1.n_cols) = tmp1;
		tmp1.each_col() += per.col(0);
		result.cols(2+tmp1.n_cols,1+tmp1.n_cols*2) = tmp1;
		tmp2.each_col() -= per.col(0);
		result.cols(2+tmp1.n_cols*2,1+tmp1.n_cols*3) = tmp2;
	}
	return result;
};

// This setter function sets the periodicity, calculates the reciprocity and brillioun zone. Note this resets the unit, see reset().
void Unit::setPeriodicity(const imat &p)
{
	removePeriodicity();

	if(p.n_cols == 0)
		return;
		
	using namespace arma;
	
	periodicity = p;
	offsets = recursive_find_p_vectors(p);

	// Now detmine the reciprocal lattice.
	setReciprocity();
	// Now the Brillouin Zone can be found.
	setBrillouin();
}
// This setter function removes the currently set periodicity. Note this resets the unit if a current periodicity has been removed, see reset().
void Unit::removePeriodicity() { periodicity.clear(); reciprocity.clear(); brillouin.clear(); reset(); }

// This setter function calculates bonds according to the current Atom::list atoms.
void Unit::initialize()
{
	reset(); 

	setBonds();

	initialized = true;
}

// This setter function resets the initalization (clears bonds);
void Unit::reset() { bonds.clear(); bondlocations.clear(); initialized = false; }

// Convenience function to transform from lattice coordinates (imat) to real space coordinates (mat).
mat Unit::toXyz(const imat &m)
{
	static const mat basischange = basisChangeMat();

	return basischange*m;
}

// Convenience function initiate the basis change matrix in toXyz().
mat Unit::basisChangeMat()
{
	double sqrtthreehalf = sqrt(3.)/2.;
	double invsqrtthree = 1./(sqrt(3.));

	mat result;

	//		i-vector:		j-vector:		h-vector:			l-vector:
	result  << 0.5			<< 0.5				<< 0.			<< 0.				<< arma::endr  // x-coordinates
			<< sqrtthreehalf<< -sqrtthreehalf	<< invsqrtthree	<< invsqrtthree		<< arma::endr  // y-coordinates
			<< 0.			<< 0.				<< 1.			<< 0.				<< arma::endr; // x-coordinates

	return result;
}

#include <algorithm>

// Convenience function that sets brillouin matrix after setting periodicity.
void Unit::setReciprocity()
{
	imat p = periodicity;
	mat xyzP = toXyz(p);

	using namespace arma;

	// The reciprocal vectors er determine by the dimensionality of the periodicity.
	const size_t dim = p.n_cols; // The vectors are placed in the columns.

	// One periodic vector means the reciprocal vector is in the same direction.
	if(dim == 1) 
	{
		// The magnetude of the reciprocal vector is inversly proportional to the periodic vetor.
		double magnsqrP = dot(xyzP.col(0),xyzP.col(0));
		reciprocity = 2*arma::datum::pi*xyzP/magnsqrP;
	} // Two periodic vectors mean that the two reciprocal vectors lie in the same plane.
	else if(dim == 2)
	{
		// Some rotations are needed for reciprocal calculation, the equivalent is to append a unit vector.
		vec3 unit("0 0 1");
		xyzP.insert_cols(2,unit);
		// And remove the corresponding reciprocal vector.
		reciprocity = 2*datum::pi*xyzP.i().t().eval().cols(0,1);
	} // Three periodic vectors gives the full set.
	else if(dim == 3)
	{
		reciprocity = 2*datum::pi*xyzP.i().t();
	}
}


// Convenience function the sets brillouin matrix after setting reciprocity.
// #TODO: This function only handles 1D periodicity and 2D when the periodicity is confined to the (x,y)-plane.
//		  Thought the functionallity is ggod endnough for now, future use may need expanding!!
void Unit::setBrillouin()
{
	// The Brillouin Zone is determined by the dimensionality of the reciprocity.
	const size_t dim = reciprocity.n_cols; // The vectors are placed in the columns.

	// Only one vector means a simple line, thus 2 points make the boundaries.
	if(dim == 1)
	{
		brillouin.col(0) = -reciprocity/2.;
		brillouin.col(1) =  reciprocity/2.;
	}// Two vectors means an area, thus 4-6 points make the boundaries.
	if(dim == 2)
	{		
		// First we need a set of closest lattice points in the reciprocal lattice space, i.e
		std::vector<vec3> latticepoints;
		latticepoints.push_back(reciprocity.col(0));
		latticepoints.push_back(reciprocity.col(1));
		latticepoints.push_back(-reciprocity.col(0));
		latticepoints.push_back(-reciprocity.col(1));
		
		// Them we sort them in order of angle. Note that in this program two vectors is the (x,y)-plane
		static const auto angle = [] (const vec3 &v) {
			double angle = atan2(v(1),v(0)); 
			return (angle > 0.) ? angle : angle + 2.*arma::datum::pi;
		};
		static const auto less_than_angle = [&] (const vec3 &v1, const vec3 &v2) {
			return angle(v1) < angle(v2);
		};
		// And Sort!
		std::sort(latticepoints.begin(), latticepoints.end(),less_than_angle);
		
		std::vector<vec3> finallatticepoints;
		finallatticepoints.reserve(8);

		// Now we where (and if) there is missing a point.
		for(size_t v1 = 0; v1 < 4; v1++)
		{
			size_t v2 = (v1 < 3) ? v1 + 1 : 0;

			// If the angle between two vectors is larger than 90 degrees, then a latticepoint is missing.
			if(abs(angle(latticepoints.at(v2)) - angle(latticepoints.at(v1))) > arma::datum::pi/2.)
			{
				finallatticepoints.push_back(latticepoints.at(v1));
				finallatticepoints.push_back(latticepoints.at(v1)+latticepoints.at(v2));
			}
			else
			{
				finallatticepoints.push_back(latticepoints.at(v1));
			}
		}

		// Now we need to sort, again!
		std::sort(finallatticepoints.begin(), finallatticepoints.end(),less_than_angle);

		// The new numer of vectors is saved for future reference.
		const size_t size = finallatticepoints.size();

		// Helper function
		static const auto get90DegreesRotated = [] (const vec3 &v) {
			vec3 result; result(0) = v(1); result(1) = -v(0);
			return result;
		};
		
		brillouin.set_size(3,size);
		brillouin.fill(arma::fill::zeros);

		// Now we can detmine the Brillouin Zone points.
		for(size_t v1 = 0; v1 < size; v1++)
		{
			size_t v2 = (v1 < size-1) ? v1 + 1 : 0;
			
			vec2 r1 = finallatticepoints[v1](arma::span(0,1));
			vec2 r2 = finallatticepoints[v2](arma::span(0,1));
			vec2 r3 = get90DegreesRotated(finallatticepoints[v1])(arma::span(0,1));
			vec2 r4 = get90DegreesRotated(finallatticepoints[v2])(arma::span(0,1));

			brillouin(arma::span(0,1),v1) = r1/2 + r3*((dot(r1,r2) - dot(r2,r2))/(2*(dot(r4,r1))));
		}
	}
}

#include <map>

struct ColumnMajorOrdering
{
	bool operator() (const arma::uvec2& x, const arma::uvec2& y) const {return x(1) < y(1) || (x(1) == y(1) && x(0) < y(0));}
};

// Convenience function that sets brillouin matrix after setting reciprocity.
void Unit::setBonds()
{
	grid();		// Generate the atom grid.

	const size_t size = atoms.size();

	if(size == 0)
		return;
	
	std::vector<std::vector<std::pair<long,Bond>>> neighbours;
	neighbours.resize(size);

	long elementsize = 0;

	for(size_t a = 0; a < size; a++)
	{
		neighbours[a].reserve(4);

		for(long i = -1; i <= 1; i++)
		{
			for(long j = -1; j <= 1; j++)
			{
				const Atom::reference::list &list = grid.getSurrounding(atoms[a],i,j);
				size_t subsize = list.size();
				for(size_t k = 0; k < subsize; k++)
				{
					Bond tmp(atoms[a],*list[k].atom,list[k].offset);

					if(tmp.isValid() 
						&& tmp.getCouplingType() != Bond::DimerNextNeighbours 
						&& tmp.getCouplingType() != Bond::NonDimerNeighbours)
					{
						neighbours[a].push_back(std::pair<long,Bond>(list[k].index,tmp));
						elementsize++;
					}
				}
			}
		}

		std::sort(neighbours[a].begin(),neighbours[a].end(), [](const std::pair<long,Bond> &v1, const std::pair<long,Bond> &v2) { return v1.first < v2.first; } );
	}
	
	bonds.clear();
	bonds.resize(elementsize);
	bondlocations.set_size(2,elementsize);
	long element = 0;
	for(size_t a = 0; a < size && element < elementsize; a++)
	{
		long nn = neighbours[a].size();
		
		for(size_t n = 0; n < nn && element < elementsize; n++)
		{
			if( n > 0 && neighbours[a][n-1].first == neighbours[a][n].first)
				element--;
			else
			{
				bondlocations(0,element) = neighbours[a][n].first;
				bondlocations(1,element) = a;
			}

			bonds[element].push_back(neighbours[a][n].second);
			element++;
		}
	}

	bonds.resize(element);
	bondlocations.resize(2,element);

	grid.clear();
}